import base64
import pickle
import json
from fastapi import APIRouter
from fastapi import Body

from service.server_app import main_service

import service.data_class as data_class
from utils import generate_timestamp_tuple


class MainController:
    def __init__(self, router: APIRouter):
        self.router = router
        self.router.add_api_route(
            "/feature_extraction/", self.feature_extraction, methods=['POST'])
        self.router.add_api_route(
            "/get_bboxes_and_scores/", self.get_bboxes_and_scores, methods=['POST'])
        self.router.add_api_route(
            "/concat_result/", self.concat_result, methods=['POST'])
        self.router.add_api_route(
            "/multiclass_nms/", self.multiclass_nms, methods=['POST'])

        self.router.add_api_route(
            "/speedtest/", self.speed_test, methods=['GET'])
        # 生成用于测速的数据
        self.__speed_test_data = {"data": "".join(str(s) for s in range(int(1e6)))}

        self.handlers = [self.__do_feature_extraction,
                         self.__do_get_bboxes_and_scores,
                         self.__do_concat_result,
                         self.__do_multiclass_nms]

    def process(self, index, last_result):
        while index < len(self.handlers):
            if last_result['action'][index] == True:
                print("Calculating {}".format(index))
                last_result = self.handlers[index](last_result)
                index += 1
            else:
                break
        return last_result

    async def speed_test(self):
        print("speed test")
        return self.__speed_test_data

    async def feature_extraction(self, last_result: data_class.last_result = Body()):

        data = json.loads(last_result.json())
        return self.process(0, data)

    # 获取预测框以及得分 参数outputs是一个ndarray的list. 函数返回的都是ndarray的list

    async def get_bboxes_and_scores(self, last_result: data_class.last_result = Body()):
        data = json.loads(last_result.json())
        return self.process(1, data)

    # 拼接不同层级的结果 total_boxes和total_scores都是ndarray的list. 函数的返回内容也都是ndarray的list

    async def concat_result(self, last_result: data_class.last_result = Body()):

        data = json.loads(last_result.json())
        return self.process(2, data)

    async def multiclass_nms(self, last_result: data_class.last_result = Body()):
        data = json.loads(last_result.json())
        return self.process(3, data)

    # 特征提取

    def __do_feature_extraction(self, last_result):
        last_result['timestamps'].append(
            generate_timestamp_tuple("feature_extraction", "begin"))
        img_data = self.b64decode(last_result['result']['img_data'])

        output = main_service.feature_extraction(img_data)

        img_scale = last_result['result']['img_scale']
        last_result['result'] = {
            'outputs': self.b64encode(output),
            "img_scale": img_scale,
        }
        last_result['timestamps'].append(
            generate_timestamp_tuple("feature_extraction", "end"))
        return last_result

    def __do_get_bboxes_and_scores(self, last_result):
        last_result['timestamps'].append(
            generate_timestamp_tuple("get_bboxes_and_scores", "begin"))
        im_shape = self.b64decode(last_result['result']['img_scale'])
        prev_step_output = self.b64decode(last_result['result']['outputs'])

        result = main_service.get_bboxes_and_scores(im_shape, prev_step_output)
        total_boxes, total_scores = result

        last_result['result'] = {
            "total_boxes": self.b64encode(total_boxes),
            "total_scores": self.b64encode(total_scores)
        }
        last_result['timestamps'].append(
            generate_timestamp_tuple("get_bboxes_and_scores", "end"))
        return last_result

    def __do_concat_result(self, last_result):
        last_result['timestamps'].append(
            generate_timestamp_tuple("concat_result", "begin"))

        total_boxes = self.b64decode(last_result['result']['total_boxes'])
        total_scores = self.b64decode(last_result['result']['total_scores'])

        result = main_service.concat_result(total_boxes, total_scores)
        yolo_boxes, yolo_scores = result

        last_result['result'] = {
            "yolo_boxes": self.b64encode(yolo_boxes),
            "yolo_scores": self.b64encode(yolo_scores)
        }
        last_result['timestamps'].append(
            generate_timestamp_tuple("concat_result", "end"))
        return last_result

    def __do_multiclass_nms(self, last_result):
        last_result['timestamps'].append(
            generate_timestamp_tuple("multiclass_nms", "begin"))

        bboxes = self.b64decode(last_result['result']['yolo_boxes'])
        scores = self.b64decode(last_result['result']['yolo_scores'])

        opt = main_service.multiclass_nms(bboxes, scores)[0]

        last_result['result'] = {
            'ret': self.b64encode(opt)
        }

        last_result['timestamps'].append(
            generate_timestamp_tuple("multiclass_nms", "end"))
        return last_result

    def b64encode(self, data):
        return str(base64.b64encode(pickle.dumps(data)), encoding='utf-8')

    def b64decode(self, data):
        return pickle.loads(base64.b64decode(bytes(data, encoding='utf8')))
